clc;
clear;

%% 迭代步骤 A 开始
% 设定罚系数
lambda = 1e12; % 罚系数

%% 迭代步骤 B 开始
% 读取网格数据
load msh

%% 迭代步骤 C 开始
% 设定物料黏度
niandu = 1000;

%% 迭代步骤 D 开始
% 设定边界条件
u1 = 0;
v1 = 0;
u3 = 0.01;
v3 = 0;
JBV1 = [BP1', u1*ones(size(BP1))', v1*ones(size(BP1))'];
JBV3 = [BP3', u3*ones(size(BP3))', v3*ones(size(BP3))'];
JBV = [JBV1; JBV3];
P2 = 0;
P4 = 1000;
JBP2 = [BE2, ones(size(BE2(:,1)))*P2, ones(size(BE2(:,1)))*P2];
JBP4 = [BE4, ones(size(BE4(:,1)))*P4, ones(size(BE4(:,1)))*P4];
JBP = [JBP2; JBP4];
clear JBV1 JBV3 BP1 BP3 BP4;
clear JBP2 JBP4 P2 P4;
clear BE1 BE2 BE3 BE4;
clear u1 v1 u3 v3;

%% 迭代步骤 E 开始
% 初始化总体方程各个数据子块
DP11 = zeros(Nz, Nz);
DP12 = zeros(Nz, Nz);
DP21 = zeros(Nz, Nz);
DP22 = zeros(Nz, Nz);
F1 = zeros(Nz, 1);
F2 = zeros(Nz, 1);

%% 迭代步骤 F 开始
% 计算单元方程系数矩阵各子块并组装
for i = 1:E
    for ie = 1:9
        JXYe(ie,:) = JXYV(JMV(i,ie),:);
    end
    [DPe11,DPe12,DPe21,DPe22] = function_of_DPe(JXYe, niandu, lambda);
    for r = 1:9
        for s = 1:9
            DP11(JMV(i,r),JMV(i,s)) = DP11(JMV(i,r),JMV(i,s)) + DPe11(r,s);
            DP12(JMV(i,r),JMV(i,s)) = DP12(JMV(i,r),JMV(i,s)) + DPe12(r,s);
            DP21(JMV(i,r),JMV(i,s)) = DP21(JMV(i,r),JMV(i,s)) + DPe21(r,s);
            DP22(JMV(i,r),JMV(i,s)) = DP22(JMV(i,r),JMV(i,s)) + DPe22(r,s);
        end
    end
end

%% 迭代步骤 G 开始
% 计算单元方程右边向量子块并组装
for i = 1:length(JBP(:,1))
    for ie = 1:9
        JXYe(ie,:) = JXYV(JMV(JBP(i,1),ie),:);
    end
    [Fe1,Fe2] = function_of_Fe(JXYe, JBP(i,:));
    for r = 1:9
        F1(JMV(JBP(i,1),r),1) = F1(JMV(JBP(i,1),r),1) + Fe1(r,1);
        F2(JMV(JBP(i,1),r),1) = F2(JMV(JBP(i,1),r),1) + Fe2(r,1);
    end
end

%% 迭代步骤 H 开始
% 组合总体方程
K = [DP11 DP12;
    DP21 DP22];
B = [-F1; -F2];

%% 迭代步骤 I 开始
% 代入 JBV 数据
N_matrix = 2*Nz;
for i = 1:length(JBV(:,1))
    II = JBV(i,1);
    u = JBV(i,2);
    for J = 1:N_matrix
        B(J) = B(J) - K(J,II)*u;
    end
    K(II,:) = zeros(1,N_matrix);
    K(:,II) = zeros(N_matrix,1);
    K(II,II) = 1;
    B(II) = u;
end
for i = 1:length(JBV(:,1))
    II = Nz + JBV(i,1);
    v = JBV(i,3);
    for J = 1:N_matrix
        B(J) = B(J) - K(J,II)*v;
    end
    K(II,:) = zeros(1,N_matrix);
    K(:,II) = zeros(N_matrix,1);
    K(II,II) = 1;
    B(II) = v;
end

%% 迭代步骤 J 开始
% 求解方程
x = K\B;
ux_k_1 = x(1:Nz);
vy_k_1 = x(1+Nz:2*Nz);

%% 迭代步骤 K 开始
% 逐个计算单元内部节点压力, 并进行累加求平均
Padd = zeros(Nz,2);
for ie = 1:E
    for i = 1:9
        JXYe(i,:) = JXYV(JMV(ie,i),:);
        uve(i,1) = ux_k_1(JMV(ie,i),:);
        uve(i,2) = vy_k_1(JMV(ie,i),:);
    end
    PIE = function_PIE(JXYe,uve,lambda);
    for i = 1:9
        Padd(JMV(ie,i),1) = Padd(JMV(ie,i),1) + PIE(i,1);
        Padd(JMV(ie,i),2) = Padd(JMV(ie,i),2) + 1;
    end
end
for i = 1:Nz
    p_k_1(i,1) = Padd(i,1)/Padd(i,2);
end

%% 迭代步骤 L 开始
% 输出Tecplot后处理结果
E = E*4
Nz = Nz
data = [JXYV, ux_k_1, vy_k_1, sqrt(ux_k_1.^2 + vy_k_1.^2), p_k_1]
JMV4 = JMV_9to4(JMV)
% 清除多余变量
clear B E II J JBV JMV JMV4 JXYV K N_matrix Nd Nz;
clear data niandu p4 u v x;
% 存储结果
save result_of_n2

%% 出口速度提取
for i = 1:length(BP2)
    UB2(i,1) = ux_k_1(BP2(i),1);
end
UB2 = UB2